package com.futu.ai.utils;

import com.futu.ai.constant.AiConfig;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 智普AI调用工具类
 *
 * @author AI Assistant
 */
public class AiUtils {

    /**
     * 业务ID模版
     */
    private static final String REQUEST_ID_TEMPLATE = "futuAI-request-%s";

    /**
     * AI调用客户端
     */
    private static final ClientV4 CLIENT = new ClientV4.Builder(AiConfig.API_KEY).build();


    /**
     * 调用AI接口(同步)
     *
     * @param prompt 提示词
     * @return AI返回的答案
     */
    public static String aiCaller(String prompt) {

        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), prompt);
        messages.add(chatMessage);
        String requestId = String.format(REQUEST_ID_TEMPLATE, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(AiConfig.MODEL_NAME)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse invokeModelApiResp = CLIENT.invokeModelApi(chatCompletionRequest);
        return invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
    }

    /**
     * 调用AI接口(SSE)
     *
     * @param prompt 提示词
     * @return AI响应流
     */
    public static Flowable<ModelData> aiCallerFlow(String prompt) {
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), prompt);
        messages.add(chatMessage);
        String requestId = String.format(REQUEST_ID_TEMPLATE, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(AiConfig.MODEL_NAME)
                .stream(Boolean.TRUE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse invokeModelApiResp = CLIENT.invokeModelApi(chatCompletionRequest);
        return invokeModelApiResp.getFlowable();
    }
    
    /**
     * 调用图片识别API接口(同步)
     *
     * @param imageBase64 图片Base64编码
     * @param prompt 提示词
     * @return AI返回的答案
     */
    public static String imageRecognitionCaller(String imageBase64, String prompt) {
        List<ChatMessage> messages = new ArrayList<>();
        
        // 创建文本提示消息
        ChatMessage textMessage = new ChatMessage(ChatMessageRole.USER.value(), prompt);
        messages.add(textMessage);
        
        // 创建图片消息
        Map<String, Object> content = new HashMap<>();
        
        // 处理Base64编码的图片
        if (!imageBase64.startsWith("data:")) {
            imageBase64 = "data:image/jpeg;base64," + imageBase64;
        }
        
        // 构建图片消息类型
        List<Map<String, Object>> contentList = new ArrayList<>();
        
        // 添加文本内容
        Map<String, Object> textContent = new HashMap<>();
        textContent.put("type", "text");
        textContent.put("text", prompt);
        contentList.add(textContent);
        
        // 添加图像内容
        Map<String, Object> imageContent = new HashMap<>();
        imageContent.put("type", "image_url");
        
        Map<String, Object> imageUrl = new HashMap<>();
        imageUrl.put("url", imageBase64);
        imageContent.put("image_url", imageUrl);
        
        contentList.add(imageContent);
        
        // 创建包含图片的消息
        ChatMessage imageMessage = new ChatMessage();
        imageMessage.setRole(ChatMessageRole.USER.value());
        imageMessage.setContent(contentList);
        
        // 替换之前的消息列表，只使用图片消息
        messages.clear();
        messages.add(imageMessage);
        
        String requestId = String.format(REQUEST_ID_TEMPLATE, System.currentTimeMillis());
        
        // 创建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(AiConfig.VISION_MODEL_NAME)  // 使用视觉模型
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
                
        // 调用API
        ModelApiResponse invokeModelApiResp = CLIENT.invokeModelApi(chatCompletionRequest);
        return invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
    }
} 